"""The HLLC Riemann solver"""

from euler import ideal_gas_w2q, ideal_gas_q2w, ideal_gas_flux_q
import numpy

########################################################################
# Direct Wave speed estimates
########################################################################
class DavisEinfeldtEuler:
    r"""The Davis and Einfeldt wave speed estimates are based on the
    Roe averaged particle and sound speeds at either side of an
    interface:

    .. math::

        S_L = \tilde{u} - \tilde{a}\,\,\, , S_R = \tilde{u} - \tilde{a}

    with the Roe-averaged velocities are:

    .. math::

        \tilde{u} = \frac{\sqrt{\rho_L}u_L +
        \sqrt{\rho_R}u_R}{\sqrt{\rho_R} + \sqrt{\rho_L}}

        \tilde{a} = \left[ (\gamma -1)(\tilde{H} -
        \frac{1}{2}\tilde{u}^2)\right]^\frac{1}{2}

        \tilde{H} = \frac{\sqrt{\rho_L}H_L +
        \sqrt{\rho_R}H_R}{\sqrt{\rho_L} + \sqrt{\rho_R}}

    """
    def __init__(self, gamma=1.4):
        self.gamma = gamma
    
    def estimate_speeds(self, wl, wr, face, S):
        """Get an estimate for the maximum and minimum wave speeds

        **Parameters**:

        wl, wr : Left and right states for this interface.

        face : fvmhsc.grids.grid1d.Face1D
            The interface for which the flux is sought.

        S : array
            Output array for the wave speed estimates

        """
        gamma = self.gamma

        rhol, ul, pl = wl[0], wl[1], wl[2]
        rhor, ur, pr = wr[0], wr[1], wr[2]

        El = 0.5 * rhol * ul * ul + pl/(gamma - 1.0)
        Er = 0.5 * rhor * ur * ur + pr/(gamma - 1.0)

        hl = ( El + pl )/rhol
        hr = ( Er + pr )/rhor

        # compute the Roe average for the enthalpy and velocity
        havg = face.compute_roe_average(hl, hr)
        uavg = face.compute_roe_average(ul, ur)

        # average sound speeds
        cavg = numpy.sqrt( (gamma-1.0)*(havg - 0.5*uavg*uavg) )

        # min and max wave speed estimates
        S[0] = sl = uavg - cavg; S[2] = sr = uavg + cavg
        
        # estimate for the middle wave (equation 10.58, Toro)
        S[1] = (pr - pl) + rhol*ul*(sl - ul) - rhor*ur*(sr - ur)
        S[1] /= ( rhol * (sl-ul) - rhor * (sr-ur) )

class HLLCEuler:

    def __init__(self, gamma, wave_speed=0):
        self.gamma = gamma

        # set the default wave speed estimator
        self.speed_estimator = DavisEinfeldtEuler(gamma)

        # number of waves for the Riemann solution
        self.nwaves = 3

    def set_speed_estimator(self, estimator):
        """Set the wave speed estimator for the HLLC solver"""
        self.speed_estimator = estimator

    def set_eulerian_godunov_flux(self, wl, wr, face):

        # get the flux for this face
        flux = face.flux

        gamma = self.gamma

        # the wave speeds and the state vector
        S = numpy.zeros(3)
        U = numpy.zeros( (3,4) )

        # get the wave speed estimates
        self.speed_estimator.estimate_speeds( wl, wr, face, S )

        # four states for the solution
        self._compute_states(wl, wr, U, S)

        # compute the flux from the states and wave speeds
        f = self._hllc_flux(U, S)

        # store the flux
        flux[:] = f

    def compute_waf_states(self, wl, wr, face, U, S):
        """Compute the four states for the WAF flux"""
        self.speed_estimator.estimate_speeds(wl, wr, face, S)
        self._compute_states(wl, wr, U, S)

    def get_flux_jump(self, U, S, i):
        """Get the flux jump across the ith wave"""
        f = self._compute_fluxes(U, S)
        delF = f[:, i+1] - f[:, i]
        return delF

    def _compute_states(self, wl, wr, U, S):
        """Compute the four states between the three waves.

        **Parameters**

        wl, wr : array
            The two states on either side of the interface.

        U : array (3,4)
            Vector for the four states

        s : array (1,3)
            Wave speed estimates for the HLLC solver.

        """
        # get the conserved variables
        ql = ideal_gas_w2q(self.gamma, wl)
        qr = ideal_gas_w2q(self.gamma, wr)

        # left and right states are the conserved variables
        U[:, 0][:] = ql
        U[:, 3][:] = qr

        # get the star states
        self._compute_star(wl, wr, U, S)

    def _compute_star(self, wl, wr, U, S):
        """Compute the star states left and right

        **Parameters**

        wl, wr : Primitive variables on either side of the interface

        U : array (3,4)
            State vector for the Riemann solver
            
        S : array (1,3)
            Wave speed estimates for the HLLC solver


        **Notes**

        The HLLC solver computes the four states between the three
        waves emerging from the interface discontinuity. The four
        states are left, star-left, star-right and right.

        This function computes the star left and right states.

        """
        rhol, ul, pl = wl[0], wl[1], wl[2]
        rhor, ur, pr = wr[0], wr[1], wr[2]

        El = 0.5 * rhol * ul * ul + pl/(self.gamma - 1.0)
        Er = 0.5 * rhor * ur * ur + pr/(self.gamma - 1.0)

        # compute the star-left state (equation 10.33 from Toro)
        U[:, 1] = numpy.array( [1.0,

                                S[1],

                                El/rhol + (S[1]-ul)*(S[1] + pl/(rhol*(S[0]-ul)))

                                ] )

        U[:, 1] *= rhol * (S[0] - ul)/(S[0] - S[1])

        # compute the star-right state (equation 10.33 from Toro) 
        U[:, 2] = numpy.array( [1.0,

                                S[1],

                                Er/rhor + (S[1]-ur)*(S[1] + pr/(rhor*(S[2]-ur)))

                                ])

        U[:, 2] *= rhor * (S[2] - ur)/(S[2] - S[1])

    def _hllc_flux(self, U, S):
        """Return the HLLC flux (equation 10.34)"""

        # compute the flux from the four states
        f = self._compute_fluxes(U, S)

        if S[0] >= 0:
            _f = f[:, 0]

        elif S[2] <= 0:
            _f = f[:, 3]

        elif S[0] <= 0 <= S[1]:
            _f = f[:, 1]

        elif S[1] <= 0 <= S[2]:
            _f = f[:, 2]

        return _f

    def _compute_fluxes(self, U, S):
        """Compute the HLLC flux at the four states using equation (10.34)"""

        gamma = self.gamma
        f = numpy.zeros_like(U)
        
        f[:, 0] = fl = ideal_gas_flux_q( gamma, U[:, 0] )
        f[:, 3] = fr = ideal_gas_flux_q( gamma, U[:, 3] )

        f[:, 1] = fl + S[0] * ( U[:, 1] - U[:, 0] )
        f[:, 2] = fr + S[2] * ( U[:, 2] - U[:, 3] )

        return f

euler_wave_speeds = [

    ("DavisEinfeldt (default)", DavisEinfeldtEuler),

    ]
